#!/usr/bin/env python3
# -*- coding: utf-8 -*-

''' 
    date: 2021-03-21
    author: baole 

    模块功能：
        从深圳证券交易所和上海证券交易所爬取所有A股股票的代码
        深圳的代码前面加上sz，比如sz000001
        上海的代码加上sh，如sh600001
'''
import os
import sys
import requests
import time
from openpyxl import load_workbook
import pandas as pd


def query_stockcode_from_szse():
    '''从深圳证券交易所获取所有的A股代码
    '''

    filepath = "output/szse.xlsx"

    # 获取A股的excel文件

    r = query_content_from_url(
        "http://www.szse.cn/api/report/ShowReport?SHOWTYPE=xlsx&CATALOGID=1110&TABKEY=tab1")
    with open(filepath, 'wb') as f:
        f.write(r.content)

    # 获取workbook对象，再获取当前的sheet，从第二行开始遍历
    # 获取第五列，保存到数组中
    li_stockcode = list()
    wb = load_workbook(filepath)
    ws = wb.active
    for row in ws.iter_rows(min_row=2):
        for cell in row:
            if cell.column == 5:
                li_stockcode.append("sz" + cell.value)

    wb.close()

    return li_stockcode


def query_stockcode_from_sse():
    '''从上海证券交易所获取所有的A股代码
    '''

    filepath = "output/sse.csv"

    # 获取A股的excel文件
    r = query_content_from_url(
        "https://query.sse.com.cn/security/stock/downloadStockListFile.do?csrcCode=&stockCode=&areaName=&stockType=1", headers={"Referer": "http://www.sse.com.cn/assortment/stock/list/share/"})
    with open(filepath, 'wb') as f:
        f.write(r.content)

    # 上证代码就是一个csv文件，我只要获取第一列就可以了。
    li_stockcode = list()
    with open(filepath) as f:
        lines = f.read().strip().split("\n")
        lines.pop(0)
        for line in lines:
            li_stockcode.append("sh" + line.split("\t")[0])

    return li_stockcode

    # df = pd.read_csv(filepath, sep="\t", encoding="GBK")
    # li_stockcode = df[df.columns[0]].to_list()
    # print(li_stockcode)

    

def query_content_from_url(url, headers={}):
    # 对于失败的request,我们会尝试5次。
    attempt = 0
    while attempt < 5:
        try:
            r = requests.get(url , headers = headers)
            return r
        except Exception as e:
            print(e)
            time.sleep(60)
        attempt += 1
        

    return None


if __name__ == '__main__':

    li_stockcode = query_stockcode_from_sse()
    li_stockcode.extend(query_stockcode_from_szse())

    # 保存到csv，一行一个代码
    with open("../../common/conf/code/all_stockcode.csv", "w") as f:
        f.write("\n".join(li_stockcode))
